def sliding_window(image, window_size, stride):
h, w, = image.shape
patches = []
positions = []
for i in range(0, h - window_size + 1, stride):
for j in range(0, w - window_size + 1, stride):
patch = image[i:i + window_size, j:j + window_size]
patches.append(patch)
positions.append((i, j))
return patches, positions
def mode_image(image, encoder, cluster_model, window_size=tile_size, stride=2):
patches, positions = sliding_window(image, window_size, stride)
# Get latent vectors for all patches
latent_vector_patch = [encoder(patch.reshape(1, window_size, window_size, 1)).numpy().flatten()
for patch in tqdm(patches)]
latent_flattened = np.array(latent_vector_patch)
latent_flattened = scaler.transform(latent_flattened)
# Clustering
cluster_ids = cluster_model.fit_predict(latent_flattened)
# Initialize character map and vote storage
char_map = np.full(image.shape[:2], -1, dtype=int) # Initialize as -1 (unassigned)
vote_dict = {} # Stores votes for each pixel location
for (i, j), cluster_id in zip(positions, cluster_ids):
if cluster_id == -1:
continue # Ignore noise
for x in range(i, i + window_size):
for y in range(j, j + window_size):
if 0 <= x < image.shape[0] and 0 <= y < image.shape[1]: # Ensure within bounds
if (x, y) not in vote_dict:
vote_dict[(x, y)] = []
vote_dict[(x, y)].append(cluster_id)
# Apply majority voting per pixel
for (x, y), votes in vote_dict.items():
char_map[x, y] = Counter(votes).most_common(1)[0][0] # Most common cluster ID
return char_map
def find_medoids(latent_vectors_train, labels, sample_size=100):
unique_labels = np.unique(labels)
medoids = []
for label in tqdm(unique_labels):
# Get points belonging to the current cluster
cluster_points = latent_vectors_train[labels == label]
n_points = len(cluster_points)
# Reduce computations for large clusters by sampling
if n_points > sample_size:
sampled_indices = np.random.choice(n_points, size=sample_size, replace=False)
cluster_points = cluster_points[sampled_indices]
# Compute pairwise distances within the sampled cluster
distances = pairwise_distances(cluster_points)
# Find the medoid (point with the smallest sum of distances)
medoid_index = np.argmin(distances.sum(axis=0))
medoids.append(cluster_points[medoid_index])
return np.array(medoids)
def compareimage_choose(original, char, selected_group=None):
"""
Static matplotlib visualization with optional selection of a specific cluster/group in char map.
Parameters:
-----------
og : np.array
Original STM image.
char : np.array
Character map (clusters).
selected_group : int or None
If provided, only this group will be shown in the character map.
"""
fig, ax = plt.subplots(1, 4, figsize=(15, 5), dpi=500)
# First image: Original
ax[0].imshow(original, origin='upper', cmap='viridis')
ax[0].set_title('Original Image')
ax[1].imshow(char, cmap='RdBu')
ax[1].set_title('Character Map')
# Second image: Character map (filtered if group selected)
if selected_group is not None:
char_filtered = np.where(char == selected_group, char, np.nan) # Set non-selected groups to NaN
ax[2].imshow(original, cmap='viridis')
color = ax[2].imshow(char_filtered, cmap='RdBu', origin='upper', vmin=np.nanmin(char), vmax=np.nanmax(char))
ax[2].set_title(f'Character map (Group {selected_group})')
else:
color = ax[2].imshow(char, cmap='RdBu', origin='upper')
ax[2].set_title('Character map (All groups)')
fig.colorbar(color, ax=ax[1], orientation='vertical', fraction=0.05)
# Third image: Overlay
im = ax[3].imshow(original, origin='upper', cmap='viridis')
im_overlay = ax[3].imshow(char, alpha=0.4, cmap='RdBu', origin='upper')
fig.colorbar(im_overlay, ax=ax[3], orientation='vertical', fraction=0.05)
ax[3].set_title('Overlay')
for axs in ax:
axs.axis('off')
plt.show()
def interactive_compareimage(og, char):
unique_groups = np.unique(char[char != -1]) # Exclude noise if needed
def update(selected_group):
compareimage_choose(og, char, selected_group)
interact(update, selected_group=[None] + list(unique_groups)) # Add None for all groups
def extract_high_error_patches(original_image, model, tile_size=32, stride=8, min_dist=10, plot_maps=False, plot_defects=False):
"""
Identifies high-error regions and extracts 32x32 patches, ensuring minimal redundancy.
Inputs:
original_image: Full STM image (2D numpy array).
model: Autoencoder model for reconstruction error.
tile_size: Size of extracted patches (default: 32x32).
stride: Step size for scanning.
min_dist: Minimum allowed distance between defect centers.
Outputs:
Extracted defect patches and visualization.
"""
h, w = original_image.shape
half_tile = tile_size // 2
error_map = np.zeros((h, w)) # Initialize error map
count_map = np.zeros((h, w)) # Track how many times each pixel is covered
# Step 1: Compute reconstruction error map
for y in tqdm(range(0, h, stride)):
for x in range(0, w, stride):
y1, y2 = min(y, h - tile_size), min(y + tile_size, h)
x1, x2 = min(x, w - tile_size), min(x + tile_size, w)
patch = original_image[y1:y2, x1:x2]
if patch.shape != (tile_size, tile_size):
patch = cv2.copyMakeBorder(patch, 0, tile_size - patch.shape[0], 0, tile_size - patch.shape[1], borderType=cv2.BORDER_REFLECT)
patch_input = patch.reshape(1, tile_size, tile_size, 1)
reconstructed_patch = model.predict(patch_input, verbose=0)[0].reshape(tile_size, tile_size)
error_patch = np.abs(patch - reconstructed_patch) ** 2
error_map[y1:y2, x1:x2] += error_patch
count_map[y1:y2, x1:x2] += 1
count_map[count_map == 0] = 1 # Avoid division by zero
error_map /= count_map # Normalize error map
raw_map = error_map
# Step 2: Compute threshold dynamically
sigma = max(2, np.std(error_map) / 5)
error_map = gaussian_filter(error_map, sigma=sigma) # Apply Gaussian smoothing
error_skewness = skew(raw_map.flatten())
threshold_percentile = 97 + 2 * min(1, error_skewness) # Between 97% and 99%
threshold = np.percentile(error_map, threshold_percentile) # Adjust percentile dynamically
high_error_ratio = np.sum(error_map > threshold) / error_map.size
print(high_error_ratio)
if high_error_ratio > 0.2: # If more than 50% of pixels exceed the threshold
error_map = np.max(error_map) - error_map # Flip the error map
print("Error map flipped due to high error coverage.")
defect_mask = (error_map > threshold).astype(np.uint8)
# Step 3: Extract defect regions using scipy.ndimage.label for connected components
labeled_mask = nd_label(defect_mask)[0]
num_labels = np.max(labeled_mask)
defect_patches = []
centers = [] # Store (x, y) positions of defect centers
error_values = [] # Store corresponding error values
# Get the centroids using center_of_mass
for label_idx in range(1, num_labels + 1):
region_mask = (labeled_mask == label_idx) # Mask for each region
y, x = center_of_mass(region_mask) # Get the centroid of the region
if y<half_tile:
y=half_tile
if y>h-half_tile:
y=h-half_tile-1
if x<half_tile:
x=half_tile
if x>w-half_tile:
x=w-half_tile-1
if np.isnan(y) or np.isnan(x):
continue
y, x = int(round(y)), int(round(x)) # Round to nearest integer
# Ensure the patch stays within bounds
y1, y2 = max(0, y - half_tile), min(h, y + half_tile)
x1, x2 = max(0, x - half_tile), min(w, x + half_tile)
patch = original_image[y1:y2, x1:x2]
# Ensure patch is exactly 32x32 (pad if needed)
if patch.shape != (tile_size, tile_size):
patch = cv2.copyMakeBorder(patch,
top=max(0, half_tile - y), bottom=max(0, y + half_tile - h),
left=max(0, half_tile - x), right=max(0, x + half_tile - w),
borderType=cv2.BORDER_CONSTANT, value=0)
defect_patches.append(patch)
centers.append((x, y))
error_values.append(error_map[y, x])
# Step 4: **Filter patches using Non-Maximum Suppression (NMS)**
if min_dist > 0:
filtered_centers, filtered_patches = [], []
error_values = np.array(error_values)
# Sort defects by error intensity (highest first)
sorted_indices = np.argsort(-error_values)
selected = np.zeros(len(centers), dtype=bool)
for i in sorted_indices:
if selected[i]: # Skip if already selected
continue
x1, y1 = centers[i]
filtered_centers.append((x1, y1))
filtered_patches.append(defect_patches[i])
selected[i] = True
# Remove nearby detections
distances = np.linalg.norm(np.array(centers) - np.array([x1, y1]), axis=1)
selected[distances < min_dist] = True # Suppress nearby detections
centers, defect_patches = filtered_centers, filtered_patches # Update lists
# Step 5: **Visualizations**
if plot_maps==True:
fig, axs = plt.subplots(1, 3, figsize=(15, 5), dpi=150)
axs[0].imshow(original_image, cmap='viridis')
axs[0].set_title('Original STM Image')
axs[1].imshow(error_map, cmap='hot')
axs[1].set_title(f'Error Map (Threshold: {threshold_percentile:.2f})')
axs[2].imshow(original_image, cmap='viridis')
axs[2].set_title('Detected Defects')
for x, y in centers:
axs[1].scatter(x, y, color='green', marker='x', lw=1, s=10)
axs[2].scatter(x, y, color='black', marker='x', lw=1, s=10)
for ax in axs:
ax.axis('off')
plt.show()
# Step 6: **Plot defect patches in a grid**
if plot_defects==True:
num_patches = len(defect_patches)
grid_size = ceil(sqrt(num_patches)) # Find nearest square root
fig, axs = plt.subplots(grid_size, grid_size, figsize=(grid_size * 2, grid_size * 2))
axs = axs.flatten()
for i, patch in enumerate(defect_patches):
axs[i].imshow(patch, cmap='viridis')
axs[i].axis('off')
axs[i].set_title(f'Defect {i+1}')
# Hide unused subplots
for j in range(i+1, len(axs)):
axs[j].axis('off')
return error_map, defect_patches, centers
def cluster_and_overlay(defect_patches, defect_centers, original_image, encoder, clusterer, tile_size=32, plot_defects=False):
"""
Clusters defect patches using pre-trained encoder and KMeans, then overlays the clusters
onto the original image.
Inputs:
defect_patches: List of 32x32 defect patches.
defect_centers: List of (x, y) positions of defects.
original_image: The original STM image.
encoder: Pre-trained encoder model (used for feature extraction).
kmeans: Pre-trained KMeans model for clustering.
tile_size: Size of the defect patches (default 32).
Outputs:
Clustered defect visualization with overlay on the original image.
"""
# Step 1: Extract features from defect patches using the encoder
defect_features = encoder.predict(np.array(defect_patches)) # Use encoder to generate features
defect_features = scaler.transform(defect_features)
# Step 2: Perform clustering using the pre-trained KMeans model
cluster_labels = clusterer.fit_predict(defect_features)
# Step 3: Overlay cluster results onto the original image
plt.figure(figsize=(8, 8), dpi=100)
plt.imshow(original_image, cmap='viridis')
ax = plt.gca()
for (x, y), cluster in zip(defect_centers, cluster_labels):
rect = patches.Rectangle((x - 16, y - 16), 32, 32, linewidth=1.5, edgecolor='k', facecolor='none')
ax.add_patch(rect)
plt.text(x + 10, y + 10, str(cluster), color='k', fontsize=10, ha='center', va='center')
plt.title("Clustered Defects Overlay on STM Image")
plt.axis('off')
plt.show()
if plot_defects==True:
# Sort patches by cluster label
sorted_indices = np.argsort(cluster_labels)
sorted_patches = [defect_patches[i] for i in sorted_indices]
sorted_labels = [cluster_labels[i] for i in sorted_indices]
num_patches = len(defect_patches)
grid_size = ceil(sqrt(num_patches)) # Find nearest square root
if len(defect_patches)>1:
fig, axs = plt.subplots(grid_size, grid_size, figsize=(grid_size * 2, grid_size * 2))
axs = axs.flatten()
for i, (patch, label) in enumerate(zip(sorted_patches, sorted_labels)):
axs[i].imshow(patch, cmap='viridis', vmin=0, vmax=1)
axs[i].axis('off')
axs[i].set_title(f'Group {label}')
# Hide unused subplots
for j in range(i + 1, len(axs)):
axs[j].axis('off')
plt.tight_layout()
plt.show()
else:
plt.imshow(np.array(defect_patches).reshape(32,32))
plt.title(f'Group {cluster_labels}')
plt.show()
return cluster_labels
def generate_fine_characterization_mask(original_image, char_map, error_map, target_clusters, error_percentile=90, plot_mask=False):
"""
Generates a binary mask highlighting areas for fine characterization based on:
- Regions belonging to specific clusters in the characterization map.
- Regions with high reconstruction error in the error map.
Parameters:
original_image (np.array): The original STM image.
char_map (np.array): The coarse characterization map.
error_map (np.array): The reconstruction error map.
target_clusters (list): List of cluster labels to include.
error_percentile (float): Percentile threshold for error map (default: 95).
sigma (float): Gaussian smoothing parameter for the error map (default: 3).
plot_mask (bool): Whether to plot the resulting mask.
Returns:
np.array: Binary mask (same shape as original image) highlighting areas for fine characterization.
"""
h, w = original_image.shape
# Step 1: Identify high-error regions
error_threshold = np.percentile(error_map, error_percentile)
high_error_mask = error_map > error_threshold
# Step 2: Identify regions belonging to target clusters
cluster_mask = np.isin(char_map, target_clusters)
# Step 3: Combine masks (logical OR)
fine_characterization_mask = (high_error_mask | cluster_mask).astype(np.uint8)
# Step 4: Visualization (optional)
if plot_mask==True:
plt.figure(figsize=(15, 5))
plt.subplot(1, 4, 1)
plt.imshow(original_image)
plt.title('Original Image')
plt.axis('off')
plt.subplot(1, 4, 2)
plt.imshow(original_image)
char = plt.imshow(char_map, cmap='RdBu', alpha=0.8)
plt.colorbar(cmap=char)
plt.title('Characterization Map')
plt.axis('off')
plt.subplot(1, 4, 3)
plt.imshow(error_map, cmap='hot')
plt.title('Error Map')
plt.axis('off')
plt.subplot(1, 4, 4)
plt.imshow(original_image)
mask = plt.imshow(fine_characterization_mask, cmap='Reds', alpha=0.5)
plt.colorbar(cmap=mask)
plt.title('Fine Characterization Mask')
plt.axis('off')
plt.show()
return np.array(fine_characterization_mask)
def extract_masked_patches_and_cluster(original_image, fine_mask, encoder, clusterer, scaler, tile_size=32, min_dist=16,
plot_overlay=True, plot_defects=True):
"""
Extracts and clusters defect patches based on a fine characterization mask using encoder + clustering model.
Adds defect centers from the high_error_centers list to the clustering and patch extraction process.
"""
h, w = original_image.shape
half_tile = tile_size // 2
labeled_mask, num_labels = nd_label(fine_mask)
print(f"Number of regions detected: {num_labels}")
defect_patches = []
centers = []
# Step 1: Iterate over each labeled region and extract centered patches
stride = min_dist # Overlap = tile_size - stride = 16 px
for region in measure.regionprops(labeled_mask):
cy, cx = region.centroid
cy, cx = int(round(cy)), int(round(cx))
# Bounding box of the region
y0, x0, y1, x1 = region.bbox
region_h = y1 - y0
region_w = x1 - x0
if region_h <= tile_size and region_w <= tile_size:
# Small region: single centered patch (as before)
y1_patch = max(0, min(h - tile_size, cy - half_tile))
x1_patch = max(0, min(w - tile_size, cx - half_tile))
patch = original_image[y1_patch:y1_patch + tile_size, x1_patch:x1_patch + tile_size]
if patch.shape == (tile_size, tile_size):
defect_patches.append((patch - np.min(patch)) / (np.max(patch) - np.min(patch)))
centers.append((x1_patch + half_tile, y1_patch + half_tile))
else:
region_h = y1 - y0
region_w = x1 - x0
num_y_tiles = max(1, floor((region_h - tile_size) / stride) + 1)
num_x_tiles = max(1, floor((region_w - tile_size) / stride) + 1)
for i in range(num_y_tiles):
for j in range(num_x_tiles):
y = y0 + i * stride
x = x0 + j * stride
# Ensure patch stays within image bounds
if y + tile_size > h or x + tile_size > w:
continue
patch = original_image[y:y + tile_size, x:x + tile_size]
if patch.shape == (tile_size, tile_size):
defect_patches.append((patch - np.min(patch)) / (np.max(patch) - np.min(patch)))
centers.append((x + half_tile, y + half_tile))
# Step 3: Feature extraction via encoder
defect_array = np.array(defect_patches)
print(defect_array.shape)
if len(defect_array) < 1:
return [], [], []
defect_features = encoder.predict(defect_array)
defect_features = scaler.transform(defect_features)
# Step 4: Clustering
if len(defect_features) < optimal_k_all+1:
cluster_labels = clusterer.predict(defect_features)
return cluster_labels, defect_patches, centers
cluster_labels = clusterer.fit_predict(defect_features)
return cluster_labels, defect_patches, centers
# Step 5: Overlay visualization
if plot_overlay:
plt.figure(figsize=(8, 8), dpi=100)
plt.imshow(original_image, cmap='viridis')
plt.imshow(fine_mask, cmap='binary', alpha=0.2)
ax = plt.gca()
for (x, y), cluster in zip(centers, cluster_labels):
rect = plt.Rectangle((x - half_tile, y - half_tile), tile_size, tile_size,
linewidth=1.5, edgecolor='k', facecolor='none')
ax.add_patch(rect)
plt.text(x+16, y+16, str(cluster), color='red', fontsize=10, ha='center', va='center')
plt.title("Clustered Defects from Fine Mask")
plt.axis('off')
plt.show()
# Step 6: Plot grouped defects
if plot_defects:
sorted_idx = np.argsort(cluster_labels)
sorted_patches = [defect_patches[i] for i in sorted_idx]
sorted_labels = [cluster_labels[i] for i in sorted_idx]
num_patches = len(sorted_patches)
grid_size = ceil(sqrt(num_patches))
fig, axs = plt.subplots(grid_size, grid_size, figsize=(grid_size * 2, grid_size * 2))
axs = axs.flatten()
for i, (patch, label) in enumerate(zip(sorted_patches, sorted_labels)):
axs[i].imshow(patch, cmap='viridis', vmin=0, vmax=1)
axs[i].axis('off')
axs[i].set_title(f'Group {label}')
for j in range(i + 1, len(axs)):
axs[j].axis('off')
plt.tight_layout()
plt.show()
return cluster_labels, defect_patches, centers
def clean_and_round_blobs(mask, shrink_iterations=3, grow_iterations=5, kernel_size=5, line_threshold=1.2, sigma=0.2, plot_mask=True):
"""
Removes elongated line-like regions before growing blobs into circular shapes.
Parameters:
-----------
mask : np.array
Binary mask (2D numpy array) where blobs are labeled as 1 (foreground) and background as 0.
shrink_iterations : int
Number of erosion iterations (shrinking).
grow_iterations : int
Number of dilation iterations (expanding).
kernel_size : int
Size of the structuring element (disk-like kernel).
line_threshold : float
Minimum aspect ratio (length/width) for a region to be considered a "long line" and removed.
plot : bool
Whether to plot the original and processed mask.
Returns:
--------
new_mask : np.array
Processed mask with rounded and cleaned blobs.
"""
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size))
eroded_mask = np.array(cv2.erode(mask, kernel, iterations=shrink_iterations))
gauss_mask = np.array(gaussian_filter(eroded_mask, sigma=sigma))
grown_mask = np.array(cv2.dilate(gauss_mask, kernel, iterations=grow_iterations))
labeled_mask = nd_label(grown_mask)[0] # Label connected components
clean_mask = np.zeros_like(mask) # Create a clean mask
for region in measure.regionprops(labeled_mask):
min_row, min_col, max_row, max_col = region.bbox
height = max_row - min_row
width = max_col - min_col
aspect_ratio = max(height / (width + 1e-5), width / (height + 1e-5)) # Avoid division by zero
# Keep only circular or compact blobs (aspect ratio should be below threshold)
if aspect_ratio < line_threshold:
clean_mask[min_row:max_row, min_col:max_col] = (labeled_mask[min_row:max_row, min_col:max_col] == region.label)
if plot_mask:
fig, axes = plt.subplots(nrows=1, ncols=5, figsize=(25, 5))
# Example arrays: image1, image2, image3
axes[0].imshow(mask, cmap='gray')
axes[0].set_title('Original Image')
axes[0].axis('off')
axes[1].imshow(eroded_mask, cmap='gray')
axes[1].set_title('Eroded Image')
axes[1].axis('off')
axes[2].imshow(gauss_mask, cmap='gray')
axes[2].set_title('Gaussed Image')
axes[2].axis('off')
axes[3].imshow(grown_mask, cmap='gray')
axes[3].set_title('Grown Image')
axes[3].axis('off')
axes[4].imshow(clean_mask, cmap='gray')
axes[4].set_title('Lines removed')
axes[4].axis('off')
return grown_mask
def clean_defect_mask(binary_mask, min_area=75, kernel_size=3):
"""
Cleans a binary defect mask by:
- Removing small connected components below a specified area
- Applying morphological opening to smooth the result
Parameters:
binary_mask (np.ndarray): 2D binary image (0/1 or bool) of defect mask
min_area (int): Minimum area (in pixels) for components to be kept
kernel_size (int): Size of structuring element for morphological opening
Returns:
np.ndarray: Cleaned binary mask
"""
# Ensure binary input
binary_mask = (binary_mask > 0).astype(np.uint8)
# Step 1: Connected component filtering
labels = measure.label(binary_mask, connectivity=2)
cleaned_mask = np.zeros_like(binary_mask)
for region in measure.regionprops(labels):
if region.area >= min_area:
cleaned_mask[labels == region.label] = 1
# Step 2: Morphological opening to remove thin noise
kernel = np.ones((kernel_size, kernel_size), np.uint8)
cleaned_mask = cv2.morphologyEx(cleaned_mask, cv2.MORPH_OPEN, kernel)
return cleaned_mask
# Now pass these high_error_centers into the extract_masked_patches_and_cluster function
